import base64
import os
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from typing import Optional

class SecurityManager:
    def __init__(self, key_file: str = "key.key"):
        self.key_file = key_file
        self._key = None
        self._load_or_create_key()
    
    def _load_or_create_key(self) -> None:
        if os.path.exists(self.key_file):
            try:
                with open(self.key_file, 'rb') as f:
                    self._key = f.read()
            except Exception:
                self._create_new_key()
        else:
            self._create_new_key()
    
    def _create_new_key(self) -> None:
        self._key = Fernet.generate_key()
        try:
            with open(self.key_file, 'wb') as f:
                f.write(self._key)
        except Exception as e:
            print(f"创建密钥文件失败: {e}")
    
    def encrypt_password(self, password: str) -> str:
        if not password:
            return ""
        
        try:
            fernet = Fernet(self._key)
            encrypted_password = fernet.encrypt(password.encode())
            return base64.urlsafe_b64encode(encrypted_password).decode()
        except Exception as e:
            print(f"密码加密失败: {e}")
            return ""
    
    def decrypt_password(self, encrypted_password: str) -> str:
        if not encrypted_password:
            return ""
        
        try:
            fernet = Fernet(self._key)
            encrypted_data = base64.urlsafe_b64decode(encrypted_password.encode())
            decrypted_password = fernet.decrypt(encrypted_data)
            return decrypted_password.decode()
        except Exception as e:
            print(f"密码解密失败: {e}")
            return ""
    
    def generate_master_key_from_password(self, password: str, salt: bytes = None) -> bytes:
        if salt is None:
            salt = os.urandom(16)
        
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key
    
    def is_key_valid(self) -> bool:
        try:
            test_data = "test"
            encrypted = self.encrypt_password(test_data)
            decrypted = self.decrypt_password(encrypted)
            return decrypted == test_data
        except Exception:
            return False